package utils

import (
	"archive/tar"
	"archive/zip"
	"bytes"
	"compress/gzip"
	"io"
	"os"
)

func UnpackZipFile(path string) (map[string][]byte, error) {
	buff, err := os.ReadFile(path)
	if err != nil {
		return nil, err
	}
	return UnpackZipBuffer(buff)
}

func UnpackZipBuffer(buffer []byte) (map[string][]byte, error) {
	resmap := make(map[string][]byte)
	reader := &bytes.Reader{}
	reader.Reset(buffer)
	if zipfile, err := zip.NewReader(reader, reader.Size()); err == nil {
		for _, v := range zipfile.File {
			if v.FileInfo().IsDir() {
				continue
			}

			file, err := v.Open()
			if err != nil {
				return nil, err
			}
			defer file.Close()

			buff := make([]byte, v.FileInfo().Size())
			num, err := io.ReadFull(file, buff)
			if err != nil {
				return nil, err
			}
			if num < len(buff) {
				resmap[v.Name] = buff[:num]
			} else {
				resmap[v.Name] = buff
			}
		}
	} else {
		reader.Reset(buffer)
		var tarfile *tar.Reader
		if gzipfile, err := gzip.NewReader(reader); err == nil {
			defer gzipfile.Close()
			tarfile = tar.NewReader(gzipfile)
		} else {
			reader.Reset(buffer)
			tarfile = tar.NewReader(reader)
		}

		for {
			file, err := tarfile.Next()
			if io.EOF == err {
				break
			}
			if err != nil {
				return nil, err
			}
			if file.FileInfo().IsDir() {
				continue
			}

			buff := make([]byte, file.Size)
			num, err := io.ReadFull(tarfile, buff)
			if err != nil {
				return nil, err
			}
			if num < len(buff) {
				resmap[file.Name] = buff[:num]
			} else {
				resmap[file.Name] = buff
			}
		}
	}
	return resmap, nil
}
